import pandas as pd
import numpy as np


def extract_dt(df):
    df['time'] = pd.to_datetime(df['time'], format='%m%d %H:%M:%S')
    df['date'] = df['time'].dt.date
    df['hour'] = df['time'].dt.hour
    df['weekday'] = df['time'].dt.weekday
    return df


def diff_op(df, cols_sp):
    def diff(a, cols_sp):
        col = a.name
        if col not in cols_sp:
            arr = a.tolist()
            b = pd.Series(arr[1:] + [arr[-1]]) - a
        else:
            b = a
        return b
    df1 = df.sort_values(by='time').reset_index().drop('index', axis=1)
    df2 = df1.apply(diff, axis=0, **{'cols_sp': cols_sp})
    return df2


train = pd.read_hdf('./input/test.h5')
train = extract_dt(train)
train['rnk'] = range(train.shape[0])
cols = ['x', 'y', 'v', 'd', 'time', 'rnk', 'ship']
feas = ['x', 'y', 'v', 'd', 'time', 'rnk']
cols_sp = ['rnk']
trn_ = train[cols].groupby('ship')[feas].apply(diff_op, cols_sp).\
    reset_index().drop('level_1', axis=1)
trn_.columns = ['ship_1', 'xd', 'yd', 'vd', 'dd', 'td', 'rnk']

trn1 = pd.merge(train, trn_, on='rnk', suffixes=['', ''])
trn1['sd'] = trn1['td'].dt.seconds
trn1['date'] = trn1['time'].dt.day
trn1['hour'] = trn1['time'].dt.hour
trn1['month'] = trn1['time'].dt.month

drop_cols = ['rnk', 'time', 'rnk', 'ship_1']
trn2 = trn1.drop(drop_cols, axis=1)


trn = trn2.copy()
trn['xd_x_yd'] = abs(trn['xd'])*abs(trn['yd'])
trn['xd_yd'] = pd.Series((trn['xd']**2 + trn['yd']**2)**0.5)
trn['yd/xd'] = trn['yd']/trn['xd']
trn['xd/yd'] = trn['xd']/trn['yd']
trn['xd/s'] = trn['xd']/trn['sd']
trn['yd/s'] = trn['yd']/trn['sd']
trn['vd/s'] = trn['vd']/trn['sd']
trn['dd/s'] = trn['dd']/trn['sd']
trn['xd_x_yd/s'] = trn['xd_x_yd']/trn['sd']
trn['xd_yd/s'] = trn['xd_yd']/trn['sd']
trn['v_x_s'] = trn['v']*trn['sd']
trn['d_x_s'] = trn['d']*trn['sd']
trn['xd_x_s'] = (trn['xd'] != 0)*trn['sd']
trn['yd_x_s'] = (trn['yd'] != 0)*trn['sd']
trn.replace({np.nan: 0, np.inf: 0, -np.inf: 0}, inplace=True)

cols = ['date', 'hour', 'weekday', 'xd',
       'yd', 'vd', 'dd', 'sd', 'month', 'xd_x_yd', 'xd_yd', 'yd/xd',
       'xd/yd', 'xd/s', 'yd/s', 'vd/s', 'dd/s', 'xd_x_yd/s', 'xd_yd/s',
       'v_x_s', 'd_x_s', 'xd_x_s', 'yd_x_s', 'ship']
trn[cols].to_hdf('input/test_d.h5', 'df', mode='w')

